use comp_agent::controller::ssh_execute_shell::{SshExecutor, ExecutionStatus};
use std::time::Duration;
use tokio;

// 重置全局执行结果存储
fn reset_executions() {
    use comp_agent::controller::ssh_execute_shell::EXECUTIONS;
    let mut executions = EXECUTIONS.lock().unwrap();
    executions.clear();
}

#[tokio::test]
async fn test_execute_exit_0() {
    reset_executions();
    
    // 创建SSH执行器
    let mut executor = create_test_executor().expect("Failed to create SSH executor");

    // 执行命令
    let execution_id = executor.execute(None, "echo 'Hello World' && exit 0").await
        .expect("Failed to execute command");

    // 等待命令执行完成并获取结果
    let result = wait_for_execution(&mut executor, &execution_id, 2).await
        .expect("Failed to get execution result");

    // 断言
    assert_eq!(result.status, ExecutionStatus::Completed);
    assert!(result.stdout.contains("Hello") || result.stdout.contains("[EXIT_CODE: 0]"));
    assert_eq!(result.exit_code, Some(0));
}

#[tokio::test]
async fn test_execute_exit_1() {
    reset_executions();
    
    // 创建SSH执行器
    let mut executor = create_test_executor().expect("Failed to create SSH executor");

    // 执行命令
    let execution_id = executor.execute(None, "echo 'Hello World' && exit 1").await
        .expect("Failed to execute command");

    // 等待命令执行完成并获取结果
    let result = wait_for_execution(&mut executor, &execution_id, 2).await
        .expect("Failed to get execution result");

    // 断言
    assert_eq!(result.status, ExecutionStatus::Completed);
    assert!(result.stdout.contains("Hello") || result.stdout.contains("[EXIT_CODE: 1]"));
    assert_eq!(result.exit_code, Some(1));
}

// 辅助函数：创建测试执行器
fn create_test_executor() -> Result<SshExecutor, Box<dyn std::error::Error>> {
    // 使用与ssh_execute_test.rs相同的测试配置
    const TEST_HOST: &str = "10.86.13.11";
    const TEST_PORT: u16 = 31222;
    const TEST_USER: &str = "app";
    const TEST_PASSWORD: &str = "app@123";
    const TIMEOUT: u64 = 10;

    println!("Connecting to {}:{} as user {}", TEST_HOST, TEST_PORT, TEST_USER);
    let executor = SshExecutor::connect(TEST_HOST, TEST_PORT, TEST_USER, TEST_PASSWORD, TIMEOUT)?;
    println!("Successfully connected to {}:{} as user {}", TEST_HOST, TEST_PORT, TEST_USER);
    Ok(executor)
}

// 辅助函数：等待并获取执行结果
async fn wait_for_execution(executor: &mut SshExecutor, execution_id: &str, wait_secs: u64) -> Result<comp_agent::controller::ssh_execute_shell::ExecutionResult, Box<dyn std::error::Error>> {
    println!("Waiting {} seconds for execution {}", wait_secs, execution_id);
    tokio::time::sleep(Duration::from_secs(wait_secs)).await;
    let result = executor.get_execution_result(execution_id).await;
    match &result {
        Ok(r) => {
            println!("Execution {} result: status={:?}, exit_code={:?}, stdout='{}', stderr='{}'",
                execution_id, r.status, r.exit_code, r.stdout, r.stderr);
        },
        Err(e) => {
            println!("Failed to get result for execution {}: {}", execution_id, e);
        }
    }
    result
}